import os
import re
from collections import defaultdict

# 路径设置
train_dir = os.path.join('data', 'NEU-DET', 'train', 'images')
valid_dir = os.path.join('data', 'NEU-DET', 'valid')

# 查找所有图片文件
def get_defect_types(directory):
    defect_types = set()
    file_counts = defaultdict(int)
    
    if os.path.exists(directory):
        for file in os.listdir(directory):
            if file.endswith('.jpg') or file.endswith('.png'):
                # 从文件名中提取缺陷类型
                match = re.match(r'([a-zA-Z_]+)_\d+', file)
                if match:
                    defect_type = match.group(1)
                    defect_types.add(defect_type)
                    file_counts[defect_type] += 1
    
    return defect_types, file_counts

# 获取训练集中的缺陷类型
train_defect_types, train_counts = get_defect_types(train_dir)

print("NEU-DET数据集缺陷类型:")
for defect_type in sorted(train_defect_types):
    print(f"- {defect_type}: {train_counts[defect_type]}张图片")

print("\n总计:")
print(f"- 缺陷类型数量: {len(train_defect_types)}")
print(f"- 总图片数量: {sum(train_counts.values())}") 